import sys
import numpy as np
from matplotlib.mlab import find

def length(a):
    """
    A numpy equivalent for MATLAB's length()
    """
    try:
        return a.shape[1]
    except IndexError:
        return a.shape[0]

def invdist(datain,xx,yy):
    """
    This function reads in a 3d gridded dataset with dimensions
    lat, lon, time with grids as specified below in lat lon
    and performs inverse distance weighting on nearest pixels to provide
    interpolated data at a point
    """
    lat = np.arange(72.5, 22.5, -2.5) # lat=72.5:-2.5:25;
    lon = np.arange(190, 297.5, 2.5) # lon=190:2.5:295;
    # % maximum of 3x3 grid of "close" pixels
    weights = np.zeros((3,3)) # weights=zeros(3,3);

    # Associate each xx/yy point with appropriate model grid
    f1 = find(abs(lon-xx)<=2.5) # f1=find(abs(lon-xx)<=2.5);
    f2 = find(abs(lat-yy)<=2.5) # f2=find(abs(lat-yy)<=2.5);

    # weights are specified for euclidean distance to pixel, no adjustment is made 
    # here for longitudinal variation w/latitude deemed negligible
    try:
        for j in range(length(f1)): # for j=1:length(f1)
            for k in range(length(f2)): # for k=1:length(f2)
                # weights(k,j)=1./sqrt((lat(f2(k))-yy).^2+(lon(f1(j))-xx).^2);
                dist = abs(lat[f2[k]]-yy)**2 + abs(lon[f1[j]]-xx)**2
                if dist > 0:
                    weights[k][j] = 1. / np.sqrt(dist)
                else:
                    raise ValueError
        # normalize by sum of weights such that it all adds to 1
        weights = weights / np.sum(weights)
    except ValueError:
        weights = np.array([0,0,0,0,1,0,0,0,0]).reshape(3,3)

    DATA = np.zeros([np.size(datain,2)]) # DATA=zeros([size(datain,3) 1]);
    for j in range(length(f1)): # for j=1:length(f1)
        for k in range(length(f2)): # for k=1:length(f2)
            if weights[k][j] > 0: # if weights(k,j)>0
                DATA = DATA + np.squeeze(datain[f2[k], f1[j], :]) * weights[k][j] # DATA=DATA+squeeze(datain(f2(k),f1(j),:)).*weights(k,j);
    return DATA # 1d array

class TextOutput(object):
    """
    A simple column text output class.
    """
    def __init__(self, print_to=sys.stdout, delim=" ", lineend="\n"):
        self.print_to = print_to
        self.delim = delim
        self.lineend = lineend
        self.series = []
        self.series_fmt = []
        self.headers = []
    def add(self,series, fmt="%d", header=""):
        """
        Add series to the eventual output
        """
        self.series.append(series)
        self.series_fmt.append(fmt)
        self.headers.append(header)
    def get_text(self):
        self.print_to.write(self.delim.join(self.headers))
        self.print_to.write(self.lineend)
        max_lenth = max(map(len,self.series))
        for i in range(max_lenth):
            line = []
            for j,s in enumerate(self.series):
                try:
                    line.append( self.series_fmt[j] % s[i] )
                except (IndexError, TypeError):
                    line.append(" ")
            self.print_to.write(self.delim.join(line))
            self.print_to.write(self.lineend)
